#include "kernel_operator.h"
using namespace AscendC;
constexpr uint32_t BUFFER_NUM = 1;
class KernelAvgpool
{
public:
    __aicore__ inline KernelAvgpool() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, uint32_t blockLength,
                                uint32_t tileLength, uint32_t lasttileLength, uint32_t tileNum,
                                uint32_t outputDim)
    {
        this->blockLength = blockLength;
        this->tileLength = tileLength / BUFFER_NUM;
        this->lasttileLength = lasttileLength;
        this->tileNum = tileNum;
        this->all_num = (int32_t)blockLength / (int32_t)outputDim;

        xGm.SetGlobalBuffer((__gm__ DTYPE_X *)x,
                            this->blockLength);
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y *)y,
                            outputDim);
        pipe.InitBuffer(xCalc, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
        pipe.InitBuffer(xCalc_last, BUFFER_NUM, this->lasttileLength * sizeof(DTYPE_X));
        pipe.InitBuffer(yCalc, this->tileLength * sizeof(DTYPE_X));
        pipe.InitBuffer(yCalc_last, this->lasttileLength * sizeof(DTYPE_X));
        pipe.InitBuffer(workSpace, this->tileLength * sizeof(DTYPE_X));
        pipe.InitBuffer(zeroCalc, this->tileLength * sizeof(DTYPE_X));
        pipe.InitBuffer(zeroCalc_last, this->lasttileLength * sizeof(DTYPE_X));
        pipe.InitBuffer(workSpace_last, this->lasttileLength * sizeof(DTYPE_X));
    }
    __aicore__ inline void Process()
    {
        uint32_t loopCount = this->tileNum;
        if (this->lasttileLength == 0)
        {
        for (int i = 0; i < loopCount; i++)
        {
            CopyIn_nl(i);
            Compute_nl(i);
        }
        }
        else
        {
            for (int i = 0; i < loopCount; i++)
            {
                CopyIn(i);
                Compute(i);
            }
        }
    }

private:
    __aicore__ inline void CopyIn(int32_t process)
    {
        LocalTensor<DTYPE_X> xLocal = xCalc.AllocTensor<DTYPE_X>();
        LocalTensor<DTYPE_X> xLocal_last = xCalc_last.AllocTensor<DTYPE_X>();
        DataCopy(xLocal, xGm[process * (this->tileLength + this->lasttileLength)], this->tileLength);
        for (int j = 0; j < this->lasttileLength; j++)
        {
            xLocal_last.SetValue(j, xGm.GetValue(process * (this->tileLength + this->lasttileLength) + this->tileLength + j));
        }
        xCalc.EnQue(xLocal);
        xCalc_last.EnQue(xLocal_last);
    }

    __aicore__ inline void CopyIn_nl(int32_t process)
    {
        LocalTensor<DTYPE_X> xLocal = xCalc.AllocTensor<DTYPE_X>();
        DataCopy(xLocal, xGm[process * (this->tileLength + this->lasttileLength)], this->tileLength);
        xCalc.EnQue(xLocal);
    }
    __aicore__ inline void Compute(int32_t process)
    {

        LocalTensor<DTYPE_X> xLocal = xCalc.DeQue<DTYPE_X>();
        LocalTensor<DTYPE_X> xLocal_last = xCalc_last.DeQue<DTYPE_X>();
        LocalTensor<DTYPE_X> yLocal = yCalc.Get<DTYPE_X>();
        LocalTensor<DTYPE_X> yLocal_last = yCalc_last.Get<DTYPE_X>();
        LocalTensor<DTYPE_X> wLocal = workSpace.Get<DTYPE_X>();
        LocalTensor<DTYPE_X> wLocal_last = workSpace_last.Get<DTYPE_X>();
        LocalTensor<DTYPE_X> zLocal = zeroCalc.Get<DTYPE_X>();
        LocalTensor<DTYPE_X> zLocal_last = zeroCalc_last.Get<DTYPE_X>();
        ReduceSum(yLocal, xLocal, wLocal, this->tileLength);
        ReduceSum(yLocal_last, xLocal_last, wLocal_last, this->lasttileLength);
        Muls(zLocal, zLocal,(DTYPE_X) 0, this->tileLength);
        Muls(zLocal_last, zLocal_last,(DTYPE_X) 0, this->tileLength);
        Adds(zLocal, zLocal, (DTYPE_X)this->all_num, this->tileLength);
        Adds(zLocal_last, zLocal_last, (DTYPE_X)this->all_num, this->lasttileLength);
        Div(yLocal, yLocal, zLocal, this->tileLength);
        Div(yLocal_last,yLocal_last,zLocal_last,this->lasttileLength);
        DTYPE_X Tvalue = ((this->tileLength > 0) ? (float)yLocal.GetValue(0) : (float)0) + (float)yLocal_last.GetValue(0);
        yGm.SetValue(process, (DTYPE_X)Tvalue);
        xCalc.FreeTensor(xLocal);
        xCalc_last.FreeTensor(xLocal_last);
        
    }

    __aicore__ inline void Compute_nl(int32_t process)
    {

        LocalTensor<DTYPE_X> xLocal = xCalc.DeQue<DTYPE_X>();
        LocalTensor<DTYPE_X> yLocal = yCalc.Get<DTYPE_X>(); 
        LocalTensor<DTYPE_X> wLocal = workSpace.Get<DTYPE_X>();
        LocalTensor<DTYPE_X> zLocal = zeroCalc.Get<DTYPE_X>();
        ReduceSum(yLocal, xLocal, wLocal, this->tileLength);
        Muls(zLocal, zLocal,(DTYPE_X) 0, this->tileLength);
        Adds(zLocal, zLocal, (DTYPE_X)this->all_num, this->tileLength);
        Div(yLocal, yLocal, zLocal, this->tileLength);
        yGm.SetValue(process, (DTYPE_X)yLocal.GetValue(0));
        xCalc.FreeTensor(xLocal);
    }

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> xCalc, xCalc_last;
    TBuf<QuePosition::VECCALC> yCalc, yCalc_last, workSpace, workSpace_last, zeroCalc, zeroCalc_last;
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<DTYPE_Y> yGm;
    uint32_t blockLength;
    uint32_t tileNum;
    uint32_t tileLength;
    uint32_t lasttileLength;
    int32_t all_num;
};
extern "C" __global__ __aicore__ void global_avg_pool(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    KernelAvgpool op;
    op.Init(x, y, tiling_data.blockLength,
            tiling_data.tileLength, tiling_data.lasttileLength, tiling_data.tileNum,
            tiling_data.outputDim);
    op.Process();
}

